import os
import subprocess
import select
import binascii
import socketserver


# 功能:
# 1. 用户通过telnet登录, 需要输入账号密码, 账号密码为code中定义的账号密码
# 2. 用户认证成功后, 相同账号的输入共享
# 3. 用户登录或者退出成功后, 相同账号登录的终端显示状态

base_any_obj = None             # 保存B侧对象
printable_words = [chr(i) for i in range(32, 127)]

# 接入码
code = {
    "b": {
        "password": "b",
        "sock_a": [],
        "sock_b": None
    },
    "a": {
        "password": "a",
        "sock_a": [],
        "sock_b": None
    },
}


# 2018.10.10 输入带颜色的文字, 注意文字需要用字符串包裹, 否则颜色不会改变.
# color: 字体颜色, info: 日志内容
def log_color(color, info):
    colors = ["gray", "red", "green", "yellow", "blue", "purple", "cyan", "white"]
    assert color in colors, "Invalid color '{}'".format(color)
    code = 30 + colors.index(color)
    return binascii.a2b_hex("1b5b306d1b5b313b33{}6d".format(code)) + info + b"\x1b\x5b\x30\x6d"


def printable(words):
    for i in range(len(words)):
        if words[i] not in printable_words:
            return False

    return True


class BaseSocket(socketserver.BaseRequestHandler):
    """
    套接字类: 主要用于对用户提供telnet服务
    作用: 该类用于处理套接字, 所有的用户套接字存储在self.socks中
    """

    def handle(self):
        global base_any_obj

        ip, port = self.client_address
        print("Connection recv from {}:{}.".format(ip, port))
        self.request.sendall(b"\xff\xfb\x01\xff\xfb\x03\xff\xfc\x1f")
        self.request.recv(1024)

        self.request.sendall(log_color("green", "Welcome {}:{}.\r\n".format(ip, port).encode("utf-8")))
        self.request.sendall(b"Please press any key to start...")
        self.request.recv(1024)
        self.request.sendall(b"\r\nlogin: ")

        while True:
            try:
                rs, ws, xs = select.select([self.request], [], [], 0.01)
                if rs:
                    data = self.request.recv(1024)
                    if not data:
                        print("Connection from {}:{} is closing.".format(ip, port))
                        raise ConnectionResetError
                    elif base_any_obj:
                        base_any_obj.write(self.request, data)
            except (ConnectionResetError, ValueError):
                print("Connection is reset by {}:{}.".format(ip, port))
                self.clear_sock(self.request)
                break

    # 删除套接字相关记录
    @classmethod
    def clear_sock(cls, request):
        global base_any_obj, code

        # 删除base_any_obj中的sock相关记录
        request_id = id(request)
        username = ""

        if base_any_obj.sock_info.get(request_id):
            username = base_any_obj.sock_info.get(request_id).get("name")
            del base_any_obj.sock_info[request_id]

        # 通知A端所有的sock连接退出的消息
        if username and code.get(username) and code.get(username).get("sock_a"):
            socks = code.get(username).get("sock_a")
            if request in socks:
                socks.pop(socks.index(request))

            for sock in socks:
                ip, port = request.getpeername()
                sock.sendall("{}:{} logout successful.\r\n".format(ip, port).encode("utf-8"))

    @classmethod
    def start(cls, host):
        thread = socketserver.ThreadingTCPServer(host, BaseSocket)
        thread.serve_forever()


class BaseAnything:
    """
    其它类: 主要用于对设备提供接口
    作用: 该类用于处理套接字, 所有的用户套接字存储再self.socks中
    A---(socket)---Agent---(Anything)---B
    """
    def __init__(self):
        self.buff = b"hello,world!"
        self.cmd = b""              # 当前命令
        self.cursor = 0             # 光标锚点
        self.sock_info = {}         # 套接字状态, {sock_id: {"status": 登录状态, "name": 用户名, "retry": 重试次数}}

    # 对接收的命令进行处理
    def process(self, request, pkt):
        print("pkt:", pkt, self.cmd, self.cursor)
        if pkt == b'\r\x00':  # Enter回车
            status, pkt = self.auth(request, self.cmd)
            if status and pkt == b"\r\n":           # 当认证成功, 且当前为回车操作, 则通过调用callback执行命令
                _, pkt = self.callback(self.cmd)
                print(f"cmd:{self.cmd},", pkt)
                if pkt:
                    pkt = b"\r\n" + pkt.encode("utf-8").replace(b"\n", b"\r\n") + b"\r\n>>> "
                else:
                    pkt = b"\r\n>>> "
            elif status:                            # 当认证成功, 且当前为输入密码操作, 则返回结果并打印标识符
                pkt += b">>> "
            self.cmd = b""
            self.cursor = 0

        elif pkt == b"\x08":  # Backspace退格
            # 只有当当光标锚点>0时, 才执行删除操作
            if self.cursor > 0:
                pkt = b"\x08\x1b\x5b\x4b" + self.cmd[self.cursor:] + b"\b" * len(self.cmd[self.cursor:])
                self.cmd = self.cmd[:self.cursor - 1] + self.cmd[self.cursor:]
                self.cursor -= 1
            else:
                pkt = b""

        elif pkt == b"\x01":  # ctrl + a跳转到命令首字符
            pkt = b"\x08" * self.cursor
            self.cursor = 0

        elif pkt == b"\x05":  # ctrl + e跳转到命令末尾
            if self.cursor != len(self.cmd):
                pkt = b"\x1b\x5b\x43" * (len(self.cmd) - self.cursor)
                self.cursor = len(self.cmd)

        elif pkt == b"\x15":  # ctrl + u清除所有命令
            pkt = b"\x08" * self.cursor + self.cmd[self.cursor:] + b"\x1b\x5b\x4b"
            pkt += b"\x08" * (len(self.cmd) - self.cursor)
            self.cmd = self.cmd[self.cursor:]
            self.cursor = 0

        elif pkt == b"\x17":  # ctrl + w清除光标前命令到上一个空格
            pre_cmd = self.cmd[:self.cursor][::-1]
            pre_cmd = pre_cmd.decode("utf-8")
            if " " in pre_cmd:
                cursor = self.cursor - pre_cmd.index(" ")
            else:
                cursor = 0

            pkt = b"\x08" * (self.cursor - cursor) + self.cmd[self.cursor:] + b"\x1b\x5b\x4b"
            pkt += b"\x08" * (len(self.cmd) - self.cursor)
            self.cmd = self.cmd[:cursor] + self.cmd[self.cursor:]
            self.cursor = cursor
            print(cursor)

        elif pkt == b'\x1b[A':  # 向上, 屏蔽命令
            pkt = b""
        elif pkt == b'\x1b[B':  # 向下, 屏蔽命令
            pkt = b""
        elif pkt == b'\x1b[C':  # 向右
            if self.cursor < len(self.cmd):
                self.cursor += 1
            else:
                pkt = b""
        elif pkt == b'\x1b[D':  # 向左
            if self.cursor:
                pkt = b"\x08"
                self.cursor -= 1
            else:
                pkt = b""

        elif pkt == b"\t":
            pkt = b""
        elif pkt == b"\x16":        # Ctrl + Shift + 6退出程序
            request.sendall(b"\r\n")
            request.close()
            return
        elif printable(pkt.decode("utf-8")):
            if self.cursor < len(self.cmd):
                self.cmd = self.cmd[:self.cursor] + pkt + self.cmd[self.cursor:]
                self.cursor += len(pkt)
                pkt += self.cmd[self.cursor:] + b"\x08" * len(self.cmd[self.cursor:])
            else:
                self.cmd += pkt
                self.cursor += len(pkt)
        else:
            print(pkt)
            pkt = b""

        return pkt

    # 对接入用户进行认证
    def auth(self, request, pkt):
        auth_status = False                                     # 认证状态
        print("auth:", pkt, self.sock_info)

        sock_info = self.sock_info.get(id(request))
        if sock_info and sock_info.get("name"):
            if self.sock_info[id(request)].get("status"):       # 认证成功
                return True, b"\r\n"

            if self.sock_info[id(request)]["retry"] >= 3:       # 超出重试次数
                self.sock_info[id(request)]["name"] = ""
                request.sendall(b"\r\nConnection has been closed by remote host.")
                request.close()
                return auth_status, b""

            if not pkt:
                pkt = b"\r\npassword: "
            else:
                username = sock_info.get("name")
                password = pkt.decode("utf-8")
                print(f"username: {username}, password: {password}")

                if code.get(username) and code.get(username).get("password") == password:
                    self.sock_info[id(request)]["status"] = 1
                    code[username]["sock_a"].append(request)
                    request.sendall(b"\r\n")
                    ip, port = request.getpeername()
                    pkt = "CLient {}:{} login successful, press Ctrl + Shift + 6 to exit.\r\n".\
                        format(ip, port).encode("utf-8")
                    auth_status = True
                else:
                    pkt = b"\r\nPermission denied, please try again.\r\nlogin: "
                    self.sock_info[id(request)]["name"] = ""
                    self.sock_info[id(request)]["retry"] += 1
        else:                                                   # 未认证/认证失败
            if pkt:
                if sock_info:
                    self.sock_info[id(request)]["name"] = pkt.decode("utf-8")
                else:
                    sock_info = dict()
                    sock_info["name"] = pkt.decode("utf-8")
                    sock_info["retry"] = 0

                    self.sock_info[id(request)] = sock_info
                pkt = b"\r\npassword: "
            else:
                pkt = b"\r\nlogin: "

        print(sock_info)

        return auth_status, pkt

    @classmethod
    def callback(cls, pkt):
        return subprocess.getstatusoutput(pkt.decode("utf-8"))

    def write(self, request, pkt):
        pkt = self.process(request, pkt)

        if pkt:
            request_id = id(request)
            if self.sock_info.get(request_id) and self.sock_info[request_id].get("status"):
                username = self.sock_info.get(request_id).get("name")
                for sock in code[username]["sock_a"]:
                    sock.sendall(pkt)
            else:
                print(request.getpeername(), pkt, self.cmd, self.cursor)
                request.sendall(pkt)

    def read(self):
        return self.buff


class Agent:
    """
    模型: A(用户侧)---(socket)---Agent---(Anything)---B(设备侧)
    功能: 代理程序用于实现代理功能, 实现用户和设备之间的双向通信
    1、代理程序监听23端口, 用于对用户提供代理服务, 套接字定义为a_socks
    2、代理程序监听1023端口, 用于处理设备侧的数据, 套接字定义为b_socks
    """
    def __init__(self):
        global base_any_obj
        base_any_obj = BaseAnything()

    @classmethod
    def start(cls):
        BaseSocket.start(("0.0.0.0", 23))


if __name__ == "__main__":
    agent = Agent()
    agent.start()
